class FactorStrategy:
    def __init__(self, data, factor_weights=None, top_n=10):
        """
        初始化策略
        :param data: 包含因子和价格数据
        :param factor_weights: 因子权重字典，如果为None则等权
        :param top_n: 持仓股票数量
        """
        self.original_data = data.copy()
        self.top_n = top_n
        self.factor_cols = [col for col in self.original_data.columns if col not in
                            ['date', 'stock', 'open', 'high', 'low', 'close',
                             'volume', 'pe', 'pb', 'market_cap', 'return']]

        if factor_weights is None:
            self.factor_weights = {factor: 1 / len(self.factor_cols) for factor in self.factor_cols}
        else:
            self.factor_weights = factor_weights

        # 预处理数据
        self._prepare_data()

    def _prepare_data(self):
        """数据预处理"""
        # 横截面标准化因子
        for factor in self.factor_cols:
            self.original_data[factor] = self.original_data.groupby('date')[factor].transform(
                lambda x: (x - x.mean()) / x.std())
        # 填充缺失值
        self.original_data = self.original_data.fillna(method='ffill').fillna(0)

    def calculate_composite_score(self, current_data=None):
        """计算综合因子得分"""
        calc_data = current_data.copy() if current_data is not None else self.original_data.copy()
        calc_data['composite_score'] = 0
        for factor, weight in self.factor_weights.items():
            if factor in calc_data.columns:
                calc_data['composite_score'] += calc_data[factor] * weight
        return calc_data

    def generate_trades(self, strategy_input):
        """
        生成交易信号
        :param strategy_input: 包含以下内容的字典:
            - current_date: 当前日期
            - current_cash: 当前现金
            - current_positions: 当前持仓 {stock: {'shares': quantity, 'entry_price': price}}
            - historical_data: 所有历史数据
            - current_portfolio_value: 当前组合总价值
        :return: 交易指令列表 [{'stock': stock, 'action': 'buy/sell', 'shares': quantity}, ...]
        """
        current_date = strategy_input['current_date']
        current_cash = strategy_input['current_cash']
        current_positions = strategy_input['current_positions']
        historical_data = strategy_input['historical_data']

        # 计算综合得分
        scored_data = self.calculate_composite_score(historical_data)
        today_data = scored_data[scored_data['date'] == current_date].copy()

        if today_data.empty:
            return []

        # 按综合得分排序
        today_data['rank'] = today_data['composite_score'].rank(ascending=False)
        top_stocks = today_data[today_data['rank'] <= self.top_n]['stock'].tolist()

        # 生成交易信号
        trades = []
        current_holdings = set(current_positions.keys())
        ideal_holdings = set(top_stocks)

        # 卖出不在top_n中的股票
        for stock in current_holdings - ideal_holdings:
            trades.append({
                'stock': stock,
                'action': 'sell',
                'shares': current_positions[stock]['shares']  # 卖出全部持仓
            })

        # 计算卖出后可用现金
        sales_proceeds = sum(
            current_positions[s]['shares'] * today_data[today_data['stock'] == s]['close'].values[0]
            for s in (current_holdings - ideal_holdings)
            if not today_data[today_data['stock'] == s].empty
        )
        available_cash = current_cash + sales_proceeds

        # 买入新的top_n股票(未持仓的)
        stocks_to_buy = ideal_holdings - current_holdings
        if stocks_to_buy and available_cash > 0:
            cash_per_stock = available_cash / len(stocks_to_buy)
            for stock in stocks_to_buy:
                stock_close = today_data[today_data['stock'] == stock]['close'].values[0]
                shares = cash_per_stock / stock_close
                trades.append({
                    'stock': stock,
                    'action': 'buy',
                    'shares': int(shares)
                })

        return trades

    def get_portfolio(self, date):
        """获取特定日期的投资组合(用于分析)
        :param date: 日期
        :return: 该日期的投资组合(股票列表)
        """
        scored_data = self.calculate_composite_score(self.original_data)
        daily_data = scored_data[scored_data['date'] == date]
        daily_data['rank'] = daily_data['composite_score'].rank(ascending=False)
        return daily_data[daily_data['rank'] <= self.top_n]['stock'].tolist()